// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/file/file_systems.h>
#include <pollux/connectors/hive/hive_connector.h>
#include <pollux/dwio/dwrf/register_dwrf_reader.h>
#include <pollux/dwio/dwrf/register_dwrf_writer.h>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/split.h>
#include <pollux/testing/exec/fuzzer_util.h>
#include <pollux/plan/input_generator.h>
#include <pollux/plan/reference_query_runner.h>
#include <pollux/testing/exec/result_verifier.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/expression/fuzzer_toolkit.h>
#include <pollux/vector/vector_fuzzer.h>
#include <pollux/vector/vector_maker.h>

TURBO_DECLARE_FLAG(int32_t, steps);

TURBO_DECLARE_FLAG(int32_t, duration_sec);

TURBO_DECLARE_FLAG(int32_t, batch_size);

TURBO_DECLARE_FLAG(int32_t, num_batches);

TURBO_DECLARE_FLAG(int32_t, max_num_varargs);

TURBO_DECLARE_FLAG(double, null_ratio);

TURBO_DECLARE_FLAG(std::string, repro_persist_path);

TURBO_DECLARE_FLAG(bool, persist_and_run_once);

TURBO_DECLARE_FLAG(bool, log_signature_stats);
using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
    using kumo::pollux::fuzzer::CallableSignature;
    using kumo::pollux::fuzzer::SignatureTemplate;

    constexpr const std::string_view kPlanNodeFileName = "plan_nodes";

    class AggregationFuzzerBase {
    public:
        AggregationFuzzerBase(
            size_t initialSeed,
            const std::unordered_map<std::string, std::shared_ptr<ResultVerifier> > &
            customVerificationFunctions,
            const std::unordered_map<std::string, std::shared_ptr<InputGenerator> > &
            customInputGenerators,
            VectorFuzzer::Options::TimestampPrecision timestampPrecision,
            const std::unordered_map<std::string, std::string> &queryConfigs,
            const std::unordered_map<std::string, std::string> &hiveConfigs,
            bool orderableGroupKeys,
            std::unique_ptr<ReferenceQueryRunner> referenceQueryRunner,
            std::optional<VectorFuzzer::Options> fuzzerOptions = std::nullopt)
            : customVerificationFunctions_{customVerificationFunctions},
              customInputGenerators_{customInputGenerators},
              queryConfigs_{queryConfigs},
              orderableGroupKeys_{orderableGroupKeys},
              persistAndRunOnce_{turbo::get_flag(FLAGS_persist_and_run_once)},
              reproPersistPath_{turbo::get_flag(FLAGS_repro_persist_path)},
              referenceQueryRunner_{std::move(referenceQueryRunner)},
              vectorFuzzer_{
                  fuzzerOptions.has_value()
                      ? fuzzerOptions.value()
                      : getFuzzerOptions(timestampPrecision),
                  pool_.get()
              } {
            filesystems::registerLocalFileSystem();
            connector::registerConnectorFactory(
                std::make_shared<connector::hive::HiveConnectorFactory>());
            registerHiveConnector(hiveConfigs);
            dwrf::registerDwrfReaderFactory();
            dwrf::registerDwrfWriterFactory();
            seed(initialSeed);
        }

        struct PlanWithSplits {
            core::PlanNodePtr plan;
            std::vector<exec::Split> splits;
        };

        struct FunctionsStats {
            size_t numFunctions = 0;
            size_t numSignatures = 0;
            size_t numSupportedFunctions = 0;
            size_t numSupportedSignatures = 0;
        };

        struct SignatureStats {
            /// Number of times a signature was chosen.
            size_t numRuns{0};

            /// Number of times generated query plan failed.
            size_t numFailed{0};
        };

        static VectorFuzzer::Options getFuzzerOptions(
            VectorFuzzer::Options::TimestampPrecision timestampPrecision) {
            VectorFuzzer::Options opts;
            opts.vectorSize = turbo::get_flag(FLAGS_batch_size);
            opts.stringVariableLength = true;
            opts.stringLength = 4'000;
            opts.nullRatio = turbo::get_flag(FLAGS_null_ratio);
            opts.timestampPrecision = timestampPrecision;
            return opts;
        }

    protected:
        struct Stats {
            // Names of functions that were tested.
            std::unordered_set<std::string> functionNames;

            // Number of iterations using aggregations over sorted inputs.
            size_t numSortedInputs{0};

            // Number of iterations where results were verified against reference DB,
            size_t numVerified{0};

            // Number of iterations where results verification was skipped because
            // function results are non-determinisic.
            size_t numVerificationSkipped{0};

            // Number of iterations where results verification was skipped because
            // reference DB doesn't support the query.
            size_t numReferenceQueryNotSupported{0};

            // Number of iterations where results verification was skipped because
            // reference DB failed to execute the query.
            size_t numReferenceQueryFailed{0};

            // Number of iterations where aggregation failed.
            size_t numFailed{0};

            void print(size_t numIterations) const;

            void updateReferenceQueryStats(ReferenceQueryErrorCode errorCode);
        };

        int32_t randInt(int32_t min, int32_t max);

        bool addSignature(
            const std::string &name,
            const FunctionSignaturePtr &signature);

        void addAggregationSignatures(
            const AggregateFunctionSignatureMap &signatureMap);

        std::shared_ptr<InputGenerator> findInputGenerator(
            const CallableSignature &signature);

        PlanWithSplits deserialize(const melon::Dynamic &obj);

        void seed(size_t seed) {
            currentSeed_ = seed;
            vectorFuzzer_.reSeed(seed);
            rng_.seed(currentSeed_);
        }

        void reSeed() {
            seed(rng_());
        }

        // Generates at least one and up to 5 scalar columns to be used as grouping,
        // partition or sorting keys.
        // Column names are generated using template '<prefix>N', where N is
        // zero-based ordinal number of the column.
        std::vector<std::string> generateKeys(
            const std::string &prefix,
            std::vector<std::string> &names,
            std::vector<TypePtr> &types);

        // Similar to generateKeys, but restricts types to orderable types (i.e. no
        // maps). For k-RANGE frame bounds, rangeFrame must be set to true so only
        // one sorting key is generated.
        std::vector<std::string> generateSortingKeys(
            const std::string &prefix,
            std::vector<std::string> &names,
            std::vector<TypePtr> &types,
            bool rangeFrame = false,
            const std::vector<TypePtr> &scalarTypes = defaultScalarTypes(),
            std::optional<uint32_t> numKeys = std::nullopt);

        std::pair<CallableSignature, SignatureStats &> pickSignature();

        std::vector<RowVectorPtr> generateInputData(
            std::vector<std::string> names,
            std::vector<TypePtr> types,
            const std::optional<CallableSignature> &signature);

        /// Generate a RowVector of the given types of children with an additional
        /// child named "row_number" of INTEGER row numbers that differentiates every
        /// row. Row numbers start from 0. This additional input vector is needed for
        /// result verification of window aggregations.
        /// @param windowFrameBounds Names of frame bound columns of a window
        /// operation. These columns are fuzzed without NULLs.
        std::vector<RowVectorPtr> generateInputDataWithRowNumber(
            std::vector<std::string> names,
            std::vector<TypePtr> types,
            const std::vector<std::string> &partitionKeys,
            const std::vector<std::string> &windowFrameBounds,
            const std::vector<std::string> &sortingKeys,
            const CallableSignature &signature);

        pollux::fuzzer::ResultOrError execute(
            const core::PlanNodePtr &plan,
            const std::vector<exec::Split> &splits = {},
            bool injectSpill = false,
            bool abandonPartial = false,
            int32_t maxDrivers = 2);

        void compare(
            const pollux::fuzzer::ResultOrError &actual,
            bool customVerification,
            const std::vector<std::shared_ptr<ResultVerifier> > &customVerifiers,
            const pollux::fuzzer::ResultOrError &expected);

        /// Returns false if the type or its children are unsupported.
        /// Currently returns false if type is Date,IntervalDayTime or Unknown.
        /// @param type
        /// @return bool
        bool isSupportedType(const TypePtr &type) const;

        // @param customVerification If false, results are compared as is. Otherwise,
        // only row counts are compared.
        // @param customVerifiers Custom verifier for each aggregate function. These
        // can be null. If not null and customVerification is true, custom verifier is
        // used to further verify the results.
        void testPlan(
            const PlanWithSplits &planWithSplits,
            bool injectSpill,
            bool abandonPartial,
            bool customVerification,
            const std::vector<std::shared_ptr<ResultVerifier> > &customVerifiers,
            const pollux::fuzzer::ResultOrError &expected,
            int32_t maxDrivers = 2);

        void printSignatureStats();

        const std::unordered_map<std::string, std::shared_ptr<ResultVerifier> >
        customVerificationFunctions_;
        const std::unordered_map<std::string, std::shared_ptr<InputGenerator> >
        customInputGenerators_;
        const std::unordered_map<std::string, std::string> queryConfigs_;

        // Whether group keys must be orderable or be just comparable.
        bool orderableGroupKeys_;
        const bool persistAndRunOnce_;
        const std::string reproPersistPath_;

        std::unique_ptr<ReferenceQueryRunner> referenceQueryRunner_;

        std::vector<CallableSignature> signatures_;
        std::vector<SignatureTemplate> signatureTemplates_;

        FunctionsStats functionsStats;

        // Stats for 'signatures_' and 'signatureTemplates_'. Stats for 'signatures_'
        // come before stats for 'signatureTemplates_'.
        std::vector<SignatureStats> signatureStats_;

        FuzzerGenerator rng_;
        size_t currentSeed_{0};

        std::shared_ptr<memory::MemoryPool> rootPool_{
            memory::memoryManager()->addRootPool()
        };
        std::shared_ptr<memory::MemoryPool> pool_{rootPool_->addLeafChild("leaf")};
        std::shared_ptr<memory::MemoryPool> writerPool_{
            rootPool_->addAggregateChild("aggregationFuzzerWriter")
        };
        VectorFuzzer vectorFuzzer_;
    };

    // Returns true if the elapsed time is greater than or equal to
    // FLAGS_duration_sec. If FLAGS_duration_sec is 0, returns true if the
    // iterations is greater than or equal to FLAGS_steps.
    template<typename T>
    bool isDone(size_t i, T startTime) {
        if (turbo::get_flag(FLAGS_duration_sec) > 0) {
            std::chrono::duration<double> elapsed =
                    std::chrono::system_clock::now() - startTime;
            return elapsed.count() >= turbo::get_flag(FLAGS_duration_sec);
        }
        return i >= turbo::get_flag(FLAGS_steps);
    }

    // Prints statistics about supported and unsupported function signatures.
    void printStats(const AggregationFuzzerBase::FunctionsStats &stats);

    // Prints (n / total) in percentage format.
    std::string printPercentageStat(size_t n, size_t total);

    // Makes an aggregation call string for the given function name and arguments.
    std::string makeFunctionCall(
        const std::string &name,
        const std::vector<std::string> &argNames,
        bool sortedInputs = false,
        bool distinctInputs = false,
        bool ignoreNulls = false);

    // Returns a list of column names from c0 to cn.
    std::vector<std::string> makeNames(size_t n);

    // Persists plans to files under basePath.
    void persistReproInfo(
        const std::vector<AggregationFuzzerBase::PlanWithSplits> &plans,
        const std::string &basePath);

    // Returns the function name used in a WindowNode. The input `node` should be a
    // pointer to a WindowNode.
    std::vector<std::string> retrieveWindowFunctionName(
        const core::PlanNodePtr &node);
} // namespace kumo::pollux::exec::test
